# 检查文件属性(exiftool)
# 搜索字符串(strings)
# 检查文件结构(binwalk)
# 尝试提取附加文件(dd)
# 分析LSB(Stegsolve)
# 尝试密码爆破(steghide)
# 检查DCT系数(jsteg)
# 手动检查Hex数据
import PIL.Image as Image
import exifread
import os
import re


class JpgSegment:
    def __init__(self, pos, name, img_bytes):
        self.pos = pos
        self.name = name
        self.img_bytes = img_bytes

    def __str__(self):
        return f"[{self.name}] (位置: 0x{self.pos:08x}, 长度: 0x{self.length:04x})"
    
    def _load_other_segments(self):
        """
        加载其他段
        """
        if self.pos + 4 > len(self.img_bytes):
            return

        self.length = int.from_bytes(
            self.img_bytes[self.pos + 2 : self.pos + 4], "big"
        )
        self.content = self.img_bytes[self.pos + 4 : self.pos + self.length]
        self.next_marker = self.pos + 2 + self.length

    def _load_sos_segment(self):
        """
        加载SOS段
        """
        next_marker = self.pos + 2
        while next_marker < len(self.img_bytes):
            if (
                self.img_bytes[next_marker] == 0xFF
                and (next_marker + 1 < len(self.img_bytes))
                and self.img_bytes[next_marker + 1] != 0x00
            ):
                break
            next_marker += 1
        self.length = next_marker - self.pos
        self.content = self.img_bytes[self.pos + 2 : next_marker]
        self.next_marker = next_marker

    def load_jpg_data(self):
        # 处理有长度字段的标记
        if self.name not in ["SOI", "EOI", "SOS"]:
            self._load_other_segments()
        elif self.name == "SOS":
            # 扫描数据直到下一个标记
            self._load_sos_segment()
        else:
            #SOI/EOI没有长度字段
            self.length = 0
            self.content = b""
            self.next_marker = self.pos + 2
    
    def print(self):
        print(
            f"位置: 0x{self.pos:08x} | 标记: {self.name}",
            end="",
        )
        print(f" | 长度: {self.length} 字节", end="")
        if self.name == "COM":
            try:
                comment = self.content.decode("utf-8", errors="replace")
                print(f" | 注释: {comment[:50]}...")
            except:
                print(" | 二进制注释内容")
        elif self.name == "APP1" and self.content.startswith(b"Exif"):
            print(" | 包含EXIF数据")
        else:
            print()


class JPGDecoder:
    def __init__(self, file_path):
        self.file_path = file_path
        self.img = None
        self.img_bytes = None
        self.is_multiple_files = False

    def _is_enable_jpg(self):
        """
        检查文件是否为JPG文件
        """
        # 检查文件是否存在
        if not os.path.exists(self.file_path):
            print("文件不存在")
            return False
        # 检查文件是否为JPG文件
        if not self.file_path.endswith(".jpg"):
            print("文件不是JPG文件")
            return False
        # 检查文件是否可读
        if not os.access(self.file_path, os.R_OK):
            print("文件不可读")
            return False
        return True

    def get_img(self):
        if self.img is not None:
            return self.img
        if self._is_enable_jpg():
            # 检查文件是否为图片文件
            try:
                self.img = Image.open(self.file_path)
                self.img.verify()
                self.img.close()
                return self.img
            except Exception as e:
                print(f"图像打开错误: {e}")
                return None

    def get_img_bytes(self):
        """
        获取图片的字节数据
        """
        if self.img_bytes is not None:
            return self.img_bytes
        with open(self.file_path, "rb") as f:
            self.img_bytes = f.read()
        return self.img_bytes

    def _show_base_info(self):
        """
        显示基本信息
        """
        print("=== 基本文件属性 ===")
        stat_info = os.stat(self.file_path)
        print(f"文件大小: {stat_info.st_size/1024:.2f} KB")

        # 使用Pillow检查
        try:
            self.get_img()
            print(f"格式: {self.img.format}")
            print(f"尺寸: {self.img.size} (宽x高)")
            print(f"模式: {self.img.mode}")
        except Exception as e:
            print(f"图像打开错误: {e}")

    def _show_exif_info(self):
        """
        显示EXIF信息
        """
        print("\n=== EXIF详细信息 ===")
        with open(self.file_path, "rb") as f:
            tags = exifread.process_file(f)
            if tags:
                for tag, value in tags.items():
                    if tag not in (
                        "JPEGThumbnail",
                        "TIFFThumbnail",
                        "Filename",
                        "EXIF MakerNote",
                    ):
                        print(f"{tag:25}: {value}")
            else:
                print("没有EXIF数据")

    def _check_embedded_files(self):
        """检查文件中是否包含多个文件"""
        # 常见文件头签名
        signatures = {
            "JPEG": b"\xff\xd8\xff",
            "PNG": b"\x89PNG",
            "GIF": b"GIF8",
            "ZIP": b"PK\x03\x04",
            "RAR": b"Rar!\x1a\x07",
            "PDF": b"%PDF",
            "BMP": b"BM",
            "ELF": b"\x7fELF",
        }
        self.get_img_bytes()
        print("===检查文件中是否包含多个文件...===")
        self.is_multiple_files = False
        for name, sig in signatures.items():
            # 查找所有匹配位置
            matches = [m.start() for m in re.finditer(re.escape(sig), self.img_bytes)]
            if matches:
                for pos in matches:
                    # 跳过文件头
                    if pos == 0:
                        continue
                    # 忽略文件头本身的签名
                    print(f"发现 {name} 文件签名在偏移量: 0x{pos:08x}")
                    self.is_multiple_files = True

    def check_file_property(self):
        """
        检查 JPEG 文件属性
        """
        self._show_base_info()
        self._show_exif_info()
        self._check_embedded_files()

    def search_string(self):
        """
        查询可能铭文存在的关键信息
        """
        # 定义常见CTF关键字模式
        patterns = [
            b"flag{.*?}",  # 字节形式的flag
            b"FLAG{.*?}",
            b"ctf{.*?}",
            b"CTF{.*?}",
            b"key{.*?}",
            b"secret{.*?}",
            b"hint{.*?}",
            b"[A-Z0-9]{10,}",  # 匹配全大写字母和数字的长字符串(可能base64)
            b"[a-f0-9]{32}",  # 匹配MD5哈希
            b"[a-f0-9]{40}",  # 匹配SHA-1哈希
            b"[a-f0-9]{64}",  # 匹配SHA-256哈希
        ]
        try:
            self.get_img_bytes()
            print(f"===搜索CTF关键字...===")
            found = False
            # 逐个模式匹配
            for pattern in patterns:
                matches = [
                    m.start() for m in re.finditer(re.escape(pattern), self.img_bytes)
                ]
                for pos in matches:
                    print(f"匹配到 [{pattern.decode()}]: (位置: 0x{pos:08x})")
                    found = True
            if not found:
                print("未找到CTF相关关键字")
        except Exception as e:
            print(f"处理文件时出错: {e}")

    def _print_file_structure(self, segments):
        """打印所有段信息"""
        for seg in segments:
            seg.print()

    def check_file_structure(self):
        print("===提取JPG文件的所有标记段===")

        segments = []

        # 已知标记名称
        marker_names = {
            "FFD8": "SOI",
            "FFE0": "APP0",
            "FFE1": "APP1",
            "FFE2": "APP2",
            "FFDB": "DQT",
            "FFC0": "SOF0",
            "FFC2": "SOF2",
            "FFC4": "DHT",
            "FFDA": "SOS",
            "FFD9": "EOI",
            "FFDD": "DRI",
            "FFFE": "COM",
        }
        pos = 0
        self.get_img_bytes()
        while pos < len(self.img_bytes):
            # 查找下一个标记(0xFF)
            marker_start = self.img_bytes.find(b"\xff", pos)
            if marker_start == -1 or marker_start + 1 >= len(self.img_bytes):
                break

            marker = self.img_bytes[marker_start : marker_start + 2]
            # 跳过填充字节(FF后跟00)
            if marker == b"\xff\x00":
                pos = marker_start + 2
                continue

            # 处理标记段
            segment = JpgSegment(
                marker_start,
                marker_names.get(marker.hex().upper(), "未知"),
                self.img_bytes,
            )
            segment.load_jpg_data()
            segments.append(segment)
            pos = segment.next_marker

        self._print_file_structure(segments)

    def extract_additional_file(self):
        pass

    def analyze_LSB(self):
        pass

    def brute_force_password(self):
        pass

    def check_DCT_coefficients(self):
        pass

    def manual_check_hex_data(self):
        pass

    def decode_analysis(self):
        self.check_file_property()
        self.search_string()
        if self.is_multiple_files:
            self.extract_additional_file()
        else:
            self.analyze_LSB()
            self.brute_force_password()


if __name__ == "__main__":
    file_path = "test-data/file_tu.jpg"
    decoder = JPGDecoder(file_path)
    # decoder.decode_analysis()
    decoder.check_file_structure()
